#if defined(_MSC_VER)
#pragma once
#endif

#ifndef LM_CORE_MEMORY_H
#define LM_CORE_MEMORY_H

#include "lm.h"
#include <windows.h>

// Synchronization Functions of MS
// http://msdn.microsoft.com/en-us/library/ms686360(v=VS.85).aspx

class ReferenceCounted {
public:
	ReferenceCounted() { nReference = 0; }
	// Use volatile long for atomic operation
	volatile LONG nReference;	
private:
	// Disable copy & assignment operator
	ReferenceCounted(const ReferenceCounted &);
	ReferenceCounted &operator=(const ReferenceCounted &);
};

// Reference class extended from pbrt
template <typename T> class Reference {
public:
	Reference(T *p = NULL) {
		ptr = p;
		if(ptr) InterlockedIncrement(&((ReferenceCounted *)ptr)->nReference);
	}
	Reference(const Reference<T> &r) {
		ptr = r.ptr;
		if(ptr) InterlockedIncrement(&((ReferenceCounted *)ptr)->nReference);
	}
	Reference &operator=(const Reference<T> &r) {
		if(ptr == r.ptr)
			return *this;
		if(r.ptr) InterlockedIncrement(&((ReferenceCounted *)r.ptr)->nReference);
		if((ptr) && 0 == InterlockedDecrement(&((ReferenceCounted *)ptr)->nReference))
			delete ptr;
		ptr = r.ptr;
		return *this;
	}
	Reference &operator=(T *p) {
		if(ptr == p)
			return *this;
		if(p) InterlockedIncrement(&((ReferenceCounted *)p)->nReference);
		if((ptr) && 0 == InterlockedDecrement(&((ReferenceCounted *)ptr)->nReference))
			delete ptr;
		ptr = p;
		return *this;
	}
	~Reference() {
		if((ptr) && 0 == InterlockedDecrement(&((ReferenceCounted *)ptr)->nReference))
			delete ptr;
	}

	inline bool operator== (const Reference &r) const { return (r.ptr == ptr); }
	inline bool operator!= (const Reference &r) const { return (r.ptr != ptr); }
	inline bool operator== (const T *p) const { return (ptr == p); }
	inline bool operator!= (const T *p) const { return (ptr != p); }

	inline T *operator->() { return ptr; }
	inline const T *operator->() const { return ptr; }
	inline T &operator*() { return *ptr; }
	inline const T &operator*() const { return *ptr; }
	inline operator T*() { return ptr; }
	inline operator bool() const { return ptr != NULL; }
	inline T *GetPtr() { return ptr; }
	inline const T *GetPtr() const { return ptr; }
private:
	T *ptr;
};

void * __restrict AllocAligned(size_t size);
template <typename T> T *AllocAligned(uint32_t count) {
	return (T *)AllocAligned(count * sizeof(T));
}
void FreeAligned(void *ptr);

// Blocked Array
template<typename T, int logBlockSize> class BlockedArray {
public:
	BlockedArray(int nu, int nv, const T *d = NULL) {
		uRes = nu;
		vRes = nv;
		uBlocks = RoundUp(uRes) >> logBlockSize;
		int nAlloc = RoundUp(uRes) * RoundUp(vRes);
		data = AllocAligned<T>(nAlloc);
		for(int i = 0; i < nAlloc; i++)
			new (&data[i]) T();
		if(d)
			for(int v = 0; v < vRes; v++)
				for(int u = 0; u < uRes; u++)
					(*this)(u,v) = d[v*uRes+u];
	}
	~BlockedArray() {
		int nAlloc = RoundUp(uRes) * RoundUp(vRes);
		for (int i = 0; i < nAlloc; i++)
			data[i].~T();
        FreeAligned(data);
	}
	T &operator()(int u, int v) {
		int bu = Block(u), bv = Block(v);
		int ou = Offset(u), ov = Offset(v);
		int offset = BlockSize()*BlockSize()*(bv*uBlocks+bu);
		offset += BlockSize()*ov+ou;
		return data[offset];
	}
	const T &operator()(int u, int v) const {
		int bu = Block(u), bv = Block(v);
		int ou = Offset(u), ov = Offset(v);
		int offset = BlockSize()*BlockSize()*(bv*uBlocks+bu);
		offset += BlockSize()*ov+ou;
		return data[offset];
	}
	void GetLinearArray(T *a) const {
		for(int v = 0; v < vRes; v++)
			for(int u = 0; u < uRes; u++)
				*a++ = (*this)(u,v);
	}
	int BlockSize() const { return 1 << logBlockSize; }
	int RoundUp(int x) const {
		return (x+BlockSize()-1) & ~(BlockSize()-1);
	}
	int uSize() const { return uRes; }
	int vSize() const { return vRes; }
	int Block(int a) { return a >> logBlockSize; }
	int Offset(int a) { return a & (BlockSize() - 1); }
private:
	T *data;
	int uRes, vRes, uBlocks;
};

// Object Arena for dynamic allocation
template<class T>
class ObjectArena {
public:
	ObjectArena() { nAvailable = 0; }
	T *Alloc() {
		// if not enough memory, reallocate more memory
		if(nAvailable == 0) {
			// Allocate 16 objects or 64kb
			int nAlloc = max((unsigned long)16,(unsigned long)(65536/sizeof(T)));
			mem = AllocAligned<T>(nAlloc);
			nAvailable = nAlloc;
			toDelete.push_back(mem);
		}
		// return the address of the memory
		nAvailable--;
		mem++;
	}
	operator T *() { return Alloc(); }
	void FreeAll() {
		for(uint32_t i = 0; i < toDelete.size(); i++)
			FreeAligned(toDelete[i]);
		toDelete.erase(toDelete.begin(), toDelete.end());
		nAvailable = 0;
	}
private:
	T *mem;
	int nAvailable;
	vector<T *> toDelete;
};

class MemoryPool {
public:
	MemoryPool(int bs = 32768) {
		blockSize = bs;
		curBlockPos = 0;
		currentBlock = AllocAligned<char>(blockSize);
	}
	~MemoryPool() {
		FreeAligned(currentBlock);
		for(uint32_t i = 0; i < usedBlocks.size(); i++)
			FreeAligned(usedBlocks[i]);
		for(uint32_t i = 0; i < availableBlocks.size(); i++)
			FreeAligned(availableBlocks[i]);
	}
	void *Alloc(int sz) {
		// Round up sz to minimum machine alignment
		sz = ((sz + 15) & (~15));
		if(curBlockPos + sz > blockSize) {
			// get new block of memory
			usedBlocks.push_back(currentBlock);
			if(availableBlocks.size() && sz <= blockSize) {
				currentBlock = availableBlocks.back();
				availableBlocks.pop_back();
			}
			else 
				currentBlock = AllocAligned<char>(max(sz, blockSize));
			curBlockPos = 0;
		}
		void *ret = currentBlock + curBlockPos;
		curBlockPos += sz;
		return ret;
	}
	template<typename T> T *Alloc(int count = 1) {
		T *ret = (T *)Alloc(count * sizeof(T));
		for(int i = 0; i < count; i++)
			new (&ret[i]) T();
		return ret;
	}
	void FreeAll() {
		curBlockPos = 0;
		while(usedBlocks.size()) {
			availableBlocks.push_back(usedBlocks.back());
			usedBlocks.pop_back();
		}
	}
private:
	int curBlockPos, blockSize;
	char *currentBlock;
	vector<char *> usedBlocks, availableBlocks;
};
#endif	// LM_CORE_MEMORY_H